package com.sun.dream.fish.secret.asymmetry;

import javax.crypto.spec.DHParameterSpec;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.InvalidParameterSpecException;
import java.util.Base64;

public class DiffieHellmanKeyPairGenerator {
    private static final int KEY_SIZE = 2048; // 调整密钥大小以满足您的要求
    private static final int EXPONENT_SIZE = 8; // 调整指数大小以满足您的要求

    public static void main(String[] args) throws NoSuchAlgorithmException, InvalidParameterSpecException, InvalidKeySpecException, InvalidAlgorithmParameterException {
        KeyPair keyPair = generateKeyPair();
        String publicKeyBase64 = encodePublicKey(keyPair.getPublic());
        String privateKeyBase64 = encodePrivateKey(keyPair.getPrivate());

        System.out.println("公钥：" + publicKeyBase64);
        System.out.println("私钥：" + privateKeyBase64);
    }

    private static KeyPair generateKeyPair() throws NoSuchAlgorithmException, InvalidParameterSpecException, InvalidKeySpecException, InvalidAlgorithmParameterException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("DH");

        AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH");
        paramGen.init(KEY_SIZE);
        AlgorithmParameters parameters = paramGen.generateParameters();

        DHParameterSpec parameterSpec = parameters.getParameterSpec(DHParameterSpec.class);

        // 使用BigInteger生成指定大小的正数指数
        BigInteger exponent = generatePositiveExponent(EXPONENT_SIZE, parameterSpec.getP());
        DHParameterSpec adjustedSpec = new DHParameterSpec(parameterSpec.getP(), parameterSpec.getG(), exponent.intValue());

        keyPairGenerator.initialize(adjustedSpec);

        return keyPairGenerator.generateKeyPair();
    }

    private static BigInteger generatePositiveExponent(int size, BigInteger modulus) {
        BigInteger exponent;
        do {
            exponent = new BigInteger(size, new SecureRandom());
        } while (exponent.signum() != 1 || exponent.compareTo(modulus) >= 0);
        return exponent;
    }

    private static String encodePublicKey(PublicKey publicKey) {
        byte[] publicKeyBytes = publicKey.getEncoded();
        return Base64.getEncoder().encodeToString(publicKeyBytes);
    }

    private static String encodePrivateKey(PrivateKey privateKey) {
        byte[] privateKeyBytes = privateKey.getEncoded();
        return Base64.getEncoder().encodeToString(privateKeyBytes);
    }
}